#! /usr/bin/env python
# coding=utf-8

import os
import shutil
from Bio import SeqIO
import sys
import argparse

# 注意 所有的序列必须有抬头 
# 读取 orthofinder的结果
# 筛选出用于构建进化树的 直系同源基因群 
# 有两个参数 present rate 是在多少比例的物种中出现
# single rate 是在多少比例的物种中为单拷贝
# 提取时 有多拷贝 选取最长的那条  


parser = argparse.ArgumentParser(
    description='''将orthofinder 得到的 直系同源群 （一般在 Orthogroup_Sequences 这个路径中） 筛选出在一定比例物种 且在 一定比例物种中为单拷贝的基因
    用法:
    filter_orthologous_groups.py -i Orthogroup_Sequences -o Orthogroup_Sequences_filtered -n 6 -p 0.5 -s 0.8 -m NIGONT,Tca
    由大天才于2021年7月12日创建于浙江农业大学''')

parser.add_argument('-i',
                help='必须给定，Orthogroup_Sequences的路径')


parser.add_argument('-o',
                help='必须给定，输出文件的路径')

parser.add_argument('-n',
                help='必须给定，物种数目')


parser.add_argument('-p',
                help='出现在多少比例的物种中，默认为0.5')


parser.add_argument('-s',
                help='在多少比例的物种中为单拷贝,默认为0.8')

parser.add_argument('-m',
                help='必须出现在物种的中的限定')

parser.add_argument('-t',action='store_true',
                help='是否删除 没有达到1：1：1的内容')





args = parser.parse_args()

if not args.i or not args.o or not args.n:
    parser.print_help()
    sys.exit()



infile_path = args.i

outfile_path  = args.o

spe_number = int(args.n)

if not args.p:
	present_rate = 0.5
else:
	present_rate = float(args.p)


if not args.s:
	single_rate = 0.8
else:
	single_rate = float(args.s)


restrain = []

if not args.m:
	pass
else:
	restrain = str(args.m).strip().split(',')

if not args.t:
	delete_double = False
else:
	delete_double = True



try:
	shutil.rmtree(outfile_path)
except:
	pass

try:
	os.mkdir(outfile_path)
except:
	pass




select_number = 0

for infile_name in os.listdir(infile_path):

	spe_presence = {}

	for s in SeqIO.parse(infile_path+'/'+infile_name,'fasta'):

		spe_name = str(s.name).split('@')[0]
		
		if spe_name not in spe_presence:
			spe_presence[spe_name] = {'count':1, 'seq':[s]}
		else:
			spe_presence[spe_name]['count'] += 1
			#if len(s.seq) > len(spe_presence[spe_name]['seq'].seq):
			spe_presence[spe_name]['seq'].append(s)

	# 物种的出现频率
	spe_presence_count = len(spe_presence.keys())
	spe_presence_single_count = sum([spe_presence[i]['count']==1 for i in spe_presence])

	spe_presence_rate = spe_presence_count/spe_number
	spe_presence_single_rate = spe_presence_single_count/spe_presence_count

	# 物种出现的限定

	spe_limit = 1

	for j in restrain:
		if j not in spe_presence:
			spe_limit = 0
			break


	
	if spe_presence_rate >= present_rate and  spe_presence_single_rate >= single_rate and spe_limit==1:
		
		outhandel = open(outfile_path+'/'+infile_name,'w')

		for i in spe_presence:

			if delete_double and spe_presence[i]['count']!=1:
				continue

			for j in spe_presence[i]['seq']:

				SeqIO.write(j,outhandel,'fasta')

		select_number += 1

		outhandel.close()



		#保存


print('筛选得到 %s 个直系同源群, 出现在 %s 的物种中，且在 %s 的物种中为单拷贝' % (select_number , present_rate,single_rate))

